package redis.inaction.ratelimiter.filter.ratelimit;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpStatus;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Mono;
import redis.inaction.ratelimiter.support.RouteApiUtils;

import java.util.Map;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setResponseStatus;

@Slf4j
public class ApiRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<ApiRateLimiterGatewayFilterFactory.Config> {

    private static final String EMPTY_KEY = "____EMPTY_KEY__";

    public static final String DEFAULT_API_CODE_NAME = "apiCode";
    public static final String DEFAULT_REPLENISH_RATE_OVERRIDE_KEY = "replenishRate";
    public static final String DEFAULT_BURST_CAPACITY_OVERRIDE_KEY = "burstCapacity";
    public static final String DEFAULT_REQUESTED_TOKENS_OVERRIDE_KEY = "requestedTokens";

    private final ApiRedisRateLimiter rateLimiter;

    private final ApiKeyResolver apiKeyResolver;

    public ApiRateLimiterGatewayFilterFactory(ApiRedisRateLimiter rateLimiter, ApiKeyResolver apiKeyResolver) {
        super(Config.class);
        this.rateLimiter = rateLimiter;
        this.apiKeyResolver = apiKeyResolver;
    }

    @Override
    public GatewayFilter apply(Config config) {
        boolean denyEmpty = getOrDefault(config.getDenyEmptyKey(), true);
        HttpStatusHolder emptyKeyStatus = HttpStatusHolder
                .parse(getOrDefault(config.getEmptyKeyStatus(), HttpStatus.FORBIDDEN.name()));

        return (exchange, chain) -> apiKeyResolver.resolve(exchange, config.apiCodeName).defaultIfEmpty(EMPTY_KEY).flatMap(apiCode -> {
            if (EMPTY_KEY.equals(apiCode)) {
                if (denyEmpty) {
                    setResponseStatus(exchange, emptyKeyStatus);
                    return exchange.getResponse().setComplete();
                }
                return chain.filter(exchange);
            }

            Mono<RateLimiter.Response> rateLimitResponse;
            Map<String, Object> api = RouteApiUtils.getRouteApi(exchange, config.apiCodeName);
            int[] rateLimitConfigs = api == null ? EMPTY_OVERRIDE_RT_CONFIG : resolveApiOverrideRateLimitConfigs(api, config);
            if( rateLimitConfigs != EMPTY_OVERRIDE_RT_CONFIG ){
                rateLimitResponse = rateLimiter.isAllowed(rateLimitConfigs[0], rateLimitConfigs[1], rateLimitConfigs[2], apiCode);
            } else {
                String routeId = config.getRouteId();
                if (routeId == null) {
                    Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                    routeId = route.getId();
                }
                rateLimitResponse = rateLimiter.isAllowed(routeId, apiCode);
            }

            return rateLimitResponse.flatMap(response -> {
                for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                    exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
                }
                if (response.isAllowed()) {
                    return chain.filter(exchange);
                }
                setResponseStatus(exchange, config.getStatusCode());
                return exchange.getResponse().setComplete();
            });
        });
    }

    private static final int[] EMPTY_OVERRIDE_RT_CONFIG = new int[0];

    private int[] resolveApiOverrideRateLimitConfigs(Map<String, Object> api, Config config){
        if(!api.containsKey(config.replenishRateOverrideKey) || !api.containsKey(config.burstCapacityOverrideKey)){
            return EMPTY_OVERRIDE_RT_CONFIG;
        }
        try{
            int replenishRate = (int)api.get(config.replenishRateOverrideKey);
            int burstCapacity = (int)api.get(config.burstCapacityOverrideKey);
            Integer requestTokens = (Integer)api.get(config.requestedTokensOverrideKey);
            return new int[]{ replenishRate, burstCapacity, requestTokens == null ? 1 : requestTokens };
        } catch (Exception e){
            log.error("resolve rate limit config override error", e);
            return EMPTY_OVERRIDE_RT_CONFIG;
        }
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return (configValue != null) ? configValue : defaultValue;
    }

    @Getter
    public static class Config extends RequestRateLimiterGatewayFilterFactory.Config {

        private String apiCodeName = DEFAULT_API_CODE_NAME;
        private String replenishRateOverrideKey = DEFAULT_REPLENISH_RATE_OVERRIDE_KEY;
        private String burstCapacityOverrideKey = DEFAULT_BURST_CAPACITY_OVERRIDE_KEY;
        private String requestedTokensOverrideKey = DEFAULT_REQUESTED_TOKENS_OVERRIDE_KEY;

        public void setApiCodeName(String apiCodeName) {
            this.apiCodeName = StringUtils.hasText(apiCodeName) ? apiCodeName : DEFAULT_API_CODE_NAME;
        }

        public void setReplenishRateOverrideKey(String replenishRateOverrideKey) {
            this.replenishRateOverrideKey = StringUtils.hasText(replenishRateOverrideKey) ? replenishRateOverrideKey : DEFAULT_REPLENISH_RATE_OVERRIDE_KEY;
        }

        public void setBurstCapacityOverrideKey(String burstCapacityOverrideKey) {
            this.burstCapacityOverrideKey = StringUtils.hasText(burstCapacityOverrideKey) ? burstCapacityOverrideKey : DEFAULT_BURST_CAPACITY_OVERRIDE_KEY;
        }

        public void setRequestedTokensOverrideKey(String requestedTokensOverrideKey) {
            this.requestedTokensOverrideKey = StringUtils.hasText(requestedTokensOverrideKey) ? requestedTokensOverrideKey : DEFAULT_REQUESTED_TOKENS_OVERRIDE_KEY;
        }
    }
}
